#!/usr/bin/env python

import cPickle as pickle
from optparse import OptionParser
import os
import os.path
import sys

import charmicat

def main():
    op = OptionParser(usage='%prog [opts] file.pkl')


    op.add_option("-f", "--force", action="store_true", dest="force", default=False, 
                  help="Force overwrite of output directory")
    op.add_option("-n", "--name", action='append', dest='names', default=[],
                  help='Write images only for histogram NAME (can be used multiple times)',
                  metavar='NAME')
    op.add_option("-o", "--output-dir", dest="output", default=None, 
                  help="Place output in directory DIR", metavar="DIR")
    op.add_option("-p", "--presentation", action="store_true", dest="presentation",
                  default=False, help="Draw images in presentation mode")
    op.add_option("-t", "--tag", dest='tag', default='',
                  help='Append TAG to all histogram titles',
                  metavar='TAG')
    op.add_option("--ext", dest='ext', default='.png', 
                  help='Write out files of type EXT', metavar='EXT')


    opts, args = op.parse_args()

    if len(args) != 1:
        sys.stderr.write("Please supply *one* pickle file with histograms.\n")
        sys.stderr.write("See `draw_hists.py -h' for more.\n")
        sys.stderr.flush()
        return -1

    infname = args[0]

    if opts.output is None:
        if infname[-4:] == 'hpkl' or infname[-4:] == '.pkl':
            opts.output = infname.rsplit('.', 1)[0]
        else:
            opts.output = 'hists'

    pngdir = os.path.join(opts.output, 'png')
    pdfdir = os.path.join(opts.output, 'pdf')

    opts.output = charmicat.safe_mkdir(opts.output, default=infname.rsplit('.',1)[0])

    fontsize = 'medium'
    if opts.presentation:
        fontsize = 'xx-large'

    if opts.ext[0] != '.':
        opts.ext = '.' + opts.ext


    infp = open(infname, 'rb')
    h    = pickle.load(infp)
    infp.close()

    if '__hatt__' in h:
        sys.path.append(os.getcwd())
        
        try:
            analysis = __import__(h['__hatt__'][0])
            h_att_path = 'analysis.' + h['__hatt__'][1]
            h_att = eval(h_att_path)
    
            for key in h:
                for attr in 'title xlabel xunits ylabel yunits zlabel zunits logz'.split():
                    if hasattr(h[key], attr) and key in h_att:
                        setattr(h[key], attr, getattr(h_att[key], attr))

            del h['__hatt__']
            
        except ImportError, e:
            sys.stderr.write("Dictionary in `{0}' has `__hatt__' attribute but module `{1}'".format(infname,
                                                                                                    h['__hatt__'][0]))
            sys.stderr.write(" specified therein cannot be imported:\n")
            sys.stderr.write(str(e) + '\n')
            sys.stderr.write("Aborting...\n")
            return -4

    if '__label__' in h:
        del h['__label__']
                            
    count = 0
    keys = h.keys()

    if len(opts.names):
        for name in opts.names:
            if name not in h:
                msg  = "Histogram `{0}' requested with `-n' option but not found in `{1}'\n"
                msg += "Aborting...\n"
                sys.stderr.write(msg.format(name, args[0]))
                return -5

        keys = opts.names

    tot   = len(keys)


    for key in keys:
        try:
            hname = os.path.join(opts.output, h[key].name) + opts.ext
            h[key].draw(name=hname, fontsize=fontsize, tag=opts.tag)
        except ValueError, e:
            sys.stderr.write("Error drawing histogram `{0}': {1}\n".format(hname, e))
            sys.stderr.write("Skipping...\n")
            sys.stderr.flush()
            continue

            
        count += 1
        sys.stdout.write('{0}/{1} complete\r'.format(count, tot))
        sys.stdout.flush()


    print 

if __name__ == "__main__":
    sys.exit(main())
